#include "batch_writer.h"
#include "exceptions.h"

#include <chrono>
#include <algorithm>

namespace kwdb {

BatchWriter::BatchWriter(const ClientConfig& config, GrpcClient& grpcClient)
    : config_(config)
    , grpcClient_(grpcClient)
    , running_(false)
    , shutdown_(false)
    , totalWritten_(0)
    , totalBatches_(0)
    , totalErrors_(0) {
}

BatchWriter::~BatchWriter() {
    stop();
}

void BatchWriter::start() {
    if (running_) {
        return;
    }
    
    running_ = true;
    shutdown_ = false;
    workerThread_ = std::thread(&BatchWriter::workerLoop, this);
}

void BatchWriter::stop() {
    if (!running_) {
        return;
    }
    
    shutdown_ = true;
    queueCondition_.notify_all();
    
    if (workerThread_.joinable()) {
        workerThread_.join();
    }
    
    running_ = false;
}

void BatchWriter::addDataPoint(const DataPoint& dataPoint) {
    if (!running_) {
        throw WriteException("批量写入器未启动");
    }
    
    {
        std::lock_guard<std::mutex> lock(queueMutex_);
        dataQueue_.push(dataPoint);
    }
    
    queueCondition_.notify_one();
}

void BatchWriter::flush() {
    if (!running_) {
        throw WriteException("批量写入器未启动");
    }
    
    std::vector<DataPoint> batch;
    {
        std::lock_guard<std::mutex> lock(queueMutex_);
        while (!dataQueue_.empty()) {
            batch.push_back(dataQueue_.front());
            dataQueue_.pop();
        }
    }
    
    if (!batch.empty()) {
        processBatch(batch);
    }
}

size_t BatchWriter::getCurrentBatchSize() const {
    std::lock_guard<std::mutex> lock(queueMutex_);
    return dataQueue_.size();
}

size_t BatchWriter::getQueueSize() const {
    std::lock_guard<std::mutex> lock(queueMutex_);
    return dataQueue_.size();
}

bool BatchWriter::isRunning() const {
    return running_;
}

void BatchWriter::workerLoop() {
    while (!shutdown_) {
        try {
            auto batch = getBatch(config_.getBatchSize(), 
                                std::chrono::milliseconds(config_.getBatchTimeout()));
            
            if (!batch.empty()) {
                processBatch(batch);
            }
        } catch (const std::exception& e) {
            totalErrors_++;
            // 记录错误日志
            // 这里可以添加日志记录
        }
    }
}

void BatchWriter::processBatch(const std::vector<DataPoint>& batch) {
    try {
        grpcClient_.writeBatch(batch);
        totalWritten_ += batch.size();
        totalBatches_++;
    } catch (const std::exception& e) {
        totalErrors_++;
        throw WriteException("批量处理失败: " + std::string(e.what()));
    }
}

std::vector<DataPoint> BatchWriter::getBatch(size_t maxSize, std::chrono::milliseconds timeout) {
    std::vector<DataPoint> batch;
    batch.reserve(maxSize);
    
    std::unique_lock<std::mutex> lock(queueMutex_);
    
    // 等待数据或超时
    if (dataQueue_.empty()) {
        queueCondition_.wait_for(lock, timeout, [this] { return !dataQueue_.empty() || shutdown_; });
    }
    
    if (shutdown_) {
        return batch;
    }
    
    // 收集批量数据
    while (!dataQueue_.empty() && batch.size() < maxSize) {
        batch.push_back(dataQueue_.front());
        dataQueue_.pop();
    }
    
    return batch;
}

bool BatchWriter::shouldFlush() const {
    return getQueueSize() >= config_.getBatchSize();
}

} // namespace kwdb
